# peppy Copyright (c) 2006-2008 Rob McMullen
# MIT Licensed; see http://peppy.flipturn.org for more info
"""
Small modifications to Ned Batchelder's Aptus fractal viewer to get it to work
inside peppy..

Note that Aptus is licensed under the MIT license, so unlike the rest of peppy
which is licensed under the GPL, this code respects the Aptus license and is
also licensed under the MIT license.
"""

import os, random, time, math
from cStringIO import StringIO

import wx
import wx.stc

from aptus.engine import AptEngine


class GaugeProgressReporter(object):
    def __init__(self, gauge=None, finishCallback=None):
        self.gauge = gauge
        self.finishCallback = finishCallback
        
    def begin(self):
        #print("starting gauge %s" % self.gauge)
        self.start = time.time()
        self.latest = self.start
        if self.gauge:
            self.gauge.startProgress("Iterating...", 100)
            wx.GetApp().cooperativeYield()

    def progress(self, frac_done, info=''):
        now = time.time()
        if now - self.latest > .25:
            #print("progess: %s" % frac_done)
            if self.gauge:
                self.gauge.updateProgress(int(100 * frac_done))
                wx.GetApp().cooperativeYield()
            self.latest = now
    
    def end(self):
        #print("ending gauge %s" % self.gauge)
        total = time.time() - self.start
        if self.gauge:
            self.gauge.stopProgress("Computed in %f seconds" % (total))
            wx.GetApp().cooperativeYield()
        if self.finishCallback:
            self.finishCallback()


class AptusMandelbrot(AptEngine):
    """ A Python wrapper around the C AptEngine class.
    """
    def __init__(self, center, diam, size, angle, iter_limit):
        self.size = size
        self.angle = angle
        
        self.pixsize = max(diam[0] / size[0], diam[1] / size[1])
        diam = self.pixsize * size[0], self.pixsize * size[1]
        
        dx = math.cos(math.radians(self.angle)) * self.pixsize
        dy = math.sin(math.radians(self.angle)) * self.pixsize

        # The upper-left corner is computed from the center, minus the radii,
        # plus half a pixel, so that we're sampling the center of the pixel.
        self.xydxdy = (dx, dy, dy, -dx)
        halfsizew = size[0]/2.0 - 0.5
        halfsizeh = size[1]/2.0 - 0.5
        self.xy0 = (
            center[0] - halfsizew * self.xydxdy[0] - halfsizeh * self.xydxdy[2],
            center[1] - halfsizew * self.xydxdy[1] - halfsizeh * self.xydxdy[3]
            )
 
        self.iter_limit = iter_limit
        self.counts = None
        self.trace_boundary = 1
        
    def coords_from_pixel(self, x, y):
        """ Get the coords of a pixel in the grid. Note that x and y can be
            fractional.
        """
        # The .5 adjustment is because the grid is aligned to the center of the
        # pixels, but we need to return the upper-left of the pixel so that other
        # math comes out right.
        x = float(x) - 0.5
        y = float(y) - 0.5
        return (
            self.xy0[0] + self.xydxdy[0]*x + self.xydxdy[2]*y,
            self.xy0[1] + self.xydxdy[1]*x + self.xydxdy[3]*y
            )

    def compute_pixels(self, progress):
        if self.counts is not None:
            return
        #print("x, y %r step %r, angle %r, iter_limit %r, size %r" % (self.xy0, self.pixsize, self.angle, self.iter_limit, self.size))

        self.clear_stats()
        progress.begin()
        
        import numpy
        self.counts = numpy.zeros((self.size[1], self.size[0]), dtype=numpy.uint32)
        self.mandelbrot_array(self.counts, progress.progress)
        progress.end()
        #print(self.get_stats())

    def color_pixels(self, palette, phase, scale=1.0):
        import numpy
        pix = numpy.zeros((self.counts.shape[0], self.counts.shape[1], 3), dtype=numpy.uint8)
        self.apply_palette(self.counts, palette.color_bytes(), phase, scale, palette.incolor, pix)
        return pix


class AptusApp(object):
    """ A mixin class for any Aptus application.
    """
    def __init__(self):
        self.center = -0.5, 0.0
        self.diam = 3.0, 3.0
        self.size = 600, 600
        self.angle = 0.0
        self.iter_limit = 999
        self.bailout = 0
        self.palette = None
        self.palette_phase = 0
        self.palette_scale = 1.0
        self.supersample = 1
        self.outfile = 'Aptus.png'
        self.continuous = False
    
    def get_fractal_class(self):
        return AptusMandelbrot
        
    def create_mandel(self):
        size = self.size[0]*self.supersample, self.size[1]*self.supersample
        fractal = self.get_fractal_class()
        m = fractal(self.center, self.diam, size, self.angle, self.iter_limit)
        # If bailout was never specified, then default differently based on
        # continuous or discrete coloring.
        if self.bailout:
            m.bailout = self.bailout
        elif self.continuous:
            m.bailout = 100.0
        else:
            m.bailout = 2.0
        if self.continuous:
            m.cont_levels = m.blend_colors = 256
        return m
    
    def color_mandel(self, m):
        return m.color_pixels(self.palette, self.palette_phase, self.palette_scale)


class AptusPanel(wx.Panel, AptusApp):
    def __init__(self, parent, mode):
        wx.Panel.__init__(self, parent, -1, style=wx.NO_BORDER+wx.WANTS_CHARS)
        AptusApp.__init__(self)
        self.mode = mode

        self.SetBackgroundStyle(wx.BG_STYLE_CUSTOM)

        self.Bind(wx.EVT_PAINT, self.on_paint)
        self.Bind(wx.EVT_LEFT_DOWN, self.on_left_down)
        self.Bind(wx.EVT_MIDDLE_DOWN, self.on_middle_down)
        self.Bind(wx.EVT_MOTION, self.on_motion)
        self.Bind(wx.EVT_LEFT_UP, self.on_left_up)
        self.Bind(wx.EVT_MIDDLE_UP, self.on_middle_up)
        self.Bind(wx.EVT_RIGHT_UP, self.on_right_up)
        self.Bind(wx.EVT_LEAVE_WINDOW, self.on_leave_window)
        self.Bind(wx.EVT_SIZE, self.on_size)
        self.Bind(wx.EVT_IDLE, self.on_idle)

        # AptusApp default values
        import aptus.palettes
        self.palette_index = 0
        self.palette = aptus.palettes.all_palettes[self.palette_index]

        self.zoom = 2.0

        self.reset_mousing()
        
        self.bitmap = None
        self.m = None
        
    def set_view(self):
        self.size = self.GetClientSize()
        self.bitmap = None

        self.m = self.create_mandel()
        self.check_size = False
        self.Refresh()

    def dilate_view(self, center, scale):
        """ Change the view by a certain scale factor, keeping the center in the
            same spot.
        """
        # Refuse to zoom out so that the whole escape circle is visible: it makes
        # boundary tracing erase the entire thing!
        if self.diam[0] * scale >= 3.9:
            return
        cx = center[0] + (self.size[0]/2 - center[0]) * scale
        cy = center[1] + (self.size[1]/2 - center[1]) * scale
        self.center = self.m.coords_from_pixel(cx, cy)
        self.diam = (self.diam[0]*scale, self.diam[1]*scale)
        self.set_view()
        
    def reset_mousing(self):
        """ Set all the mousing variables to turn rubberbanding and panning off.
        """
        self.pt_down = None
        self.rubberbanding = False
        self.rubberrect = None
        # Panning information.
        self.panning = False
        self.pt_pan = None
        self.pan_locked = False

    def finish_panning(self, mx, my):
        if not self.pt_down:
            return
        cx, cy = self.size[0]/2.0, self.size[1]/2.0
        cx -= mx - self.pt_down[0]
        cy -= my - self.pt_down[1]
        self.center = self.m.coords_from_pixel(cx, cy)
        self.set_view()
        
    def xor_rectangle(self, rect):
        dc = wx.ClientDC(self)
        dc.SetLogicalFunction(wx.XOR)
        dc.SetBrush(wx.Brush(wx.WHITE, wx.TRANSPARENT))
        dc.SetPen(wx.Pen(wx.WHITE, 1, wx.SOLID))
        dc.DrawRectangle(*rect)

    def set_cursor(self):
        # If we aren't taking input, then we shouldn't change the cursor.
        if not self.IsEnabled():
            return 
        # Set the proper cursor:
        if self.rubberbanding:
            self.SetCursor(wx.StockCursor(wx.CURSOR_MAGNIFIER))
        elif self.panning:
            self.SetCursor(wx.StockCursor(wx.CURSOR_SIZING))
        else:
            self.SetCursor(wx.StockCursor(wx.CURSOR_DEFAULT))

    # Event handlers
    
    def on_left_down(self, event):
        self.pt_down = event.GetPosition()
        self.rubberbanding = False
        if self.panning:
            self.pt_pan = self.pt_down
            self.pan_locked = False
            
    def on_middle_down(self, event):
        self.pt_down = event.GetPosition()
        self.rubberbanding = False
        self.panning = True
        self.pt_pan = self.pt_down
        self.pan_locked = False
        
    def on_motion(self, event):
        self.set_cursor()
        
        # We do nothing with mouse moves that aren't dragging.
        if not self.pt_down:
            return
        
        mx, my = event.GetPosition()
        
        if self.panning:
            if self.pt_pan != (mx, my):
                # We've moved the image: redraw it.
                self.pt_pan = (mx, my)
                self.pan_locked = True
                self.Refresh()
        else:
            if not self.rubberbanding:
                # Start rubberbanding when we have a 10-pixel rectangle at least.
                if abs(self.pt_down[0] - mx) > 10 or abs(self.pt_down[1] - my) > 10:
                    self.rubberbanding = True
    
            if self.rubberbanding:
                if self.rubberrect:
                    # Erase the old rectangle.
                    self.xor_rectangle(self.rubberrect)
                    
                self.rubberrect = (self.pt_down[0], self.pt_down[1], mx-self.pt_down[0], my-self.pt_down[1]) 
                self.xor_rectangle(self.rubberrect)
                
    def on_left_up(self, event):
        mx, my = event.GetPosition()
        if self.rubberbanding:
            # Set a new view that encloses the rectangle.
            px, py = self.pt_down
            ulx, uly = self.m.coords_from_pixel(px, py)
            lrx, lry = self.m.coords_from_pixel(mx, my)
            self.center = ((ulx+lrx)/2, (uly+lry)/2)
            self.diam = (abs(self.m.pixsize*(px-mx)), abs(self.m.pixsize*(py-my)))
            self.set_view()
        elif self.panning:
            self.finish_panning(mx, my)
        elif self.pt_down:
            # Single-click: zoom in.
            scale = self.zoom
            if event.CmdDown():
                scale = (scale - 1.0)/10 + 1.0
            self.dilate_view((mx, my), 1.0/scale)

        self.reset_mousing()        

    def on_middle_up(self, event):
        self.finish_panning(*event.GetPosition())
        self.reset_mousing()        

    def on_right_up(self, event):
        scale = self.zoom
        if event.CmdDown():
            scale = (scale - 1.0)/10 + 1.0
        self.dilate_view(event.GetPosition(), scale)
        self.reset_mousing()
        
    def on_leave_window(self, event_unused):
        if self.rubberrect:
            self.xor_rectangle(self.rubberrect)
        self.reset_mousing()
        
    def on_size(self, event_unused):
        self.check_size = True
        
    def on_idle(self, event_unused):
        self.set_cursor()
        if self.check_size and self.GetClientSize() != self.size:
            if self.GetClientSize() != (0,0):
                self.set_view()

    def on_paint(self, event_unused):
        if not self.m:
            return
        
        if not self.bitmap:
            self.bitmap = self.draw()
        
        dc = wx.AutoBufferedPaintDC(self)
        if self.panning:
            dc.SetBrush(wx.Brush(wx.Colour(128,128,128), wx.SOLID))
            dc.SetPen(wx.Pen(wx.Colour(128,128,128), 1, wx.SOLID))
            dc.DrawRectangle(0, 0, self.size[0], self.size[1])
            dc.DrawBitmap(self.bitmap, self.pt_pan[0]-self.pt_down[0], self.pt_pan[1]-self.pt_down[1], False)
        else:
            dc.DrawBitmap(self.bitmap, 0, 0, False)
            
    def draw(self):
        """ Return a bitmap with the image to display in the window.
        """
        progress = GaugeProgressReporter(self.mode.status_info, self.mode.updateText)
        self.m.compute_pixels(progress)
        pix = self.color_mandel(self.m)
        return wx.BitmapFromBuffer(pix.shape[1], pix.shape[0], pix)

    def GetClientSize(self):
        """Overridden method to use the size of the wrapper as the drawing
        area size.
        
        This method is used because of the interaction of the minibuffer with
        the main drawing area.  The minibuffer reduces the size of the drawing
        area, which causes an automatic redraw of the Aptus area.  This isn't
        desirable because it can take quite a while to calculate the new image.
        
        Using the wrapper's size means that the size used by Aptus doesn't
        change when a minibuffer is added.
        """
        return self.mode.wrapper.GetClientSize()
